{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Common Lyapunov Analysis for Linear Systems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import MathematicalProgram, Solve\n",
    "\n",
    "A = []\n",
    "if (True):\n",
    "    # Generate random stable matrices.\n",
    "    num_states = 4\n",
    "    num_systems = 2\n",
    "    for i in range(num_systems):\n",
    "        d = -np.random.rand(num_states,)\n",
    "        v = np.random.randn(num_states, num_states)\n",
    "        A.append(v.dot(np.diag(d).dot(np.linalg.inv(v))))\n",
    "else:\n",
    "    # Example from lecture notes.\n",
    "    A.append = np.array(((-1, .5), (-3, -1)))\n",
    "    A.append = np.array(((-1, .1), (-10, -1)))\n",
    "    # Interesting for 2D plotting (a two element parameterization of stable\n",
    "    # linear systems).  Stable iff ab < 1.\n",
    "    # a = randn;  ab = 2*rand - 1;  b=ab/a;\n",
    "    # A{i} = [-1 a; b -1];\n",
    "\n",
    "# Create the optimization problem.\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "# Construct an n-by-n positive semi-definite matrix as the decision\n",
    "# variables.\n",
    "num_states = A[0].shape[0]\n",
    "P = prog.NewSymmetricContinuousVariables(num_states, \"P\")\n",
    "prog.AddPositiveSemidefiniteConstraint(P - .01 * np.identity(num_states))\n",
    "\n",
    "# Add the common Lyapunov conditions.\n",
    "for i in range(len(A)):\n",
    "    # yapf: disable\n",
    "    prog.AddPositiveSemidefiniteConstraint(\n",
    "        -A[i].transpose().dot(P) - P.dot(A[i]) - .01 * np.identity(num_states))\n",
    "    # yapf: enable\n",
    "\n",
    "# Add an objective.\n",
    "prog.AddLinearCost(np.trace(P))\n",
    "\n",
    "# Run the optimization.\n",
    "result = Solve(prog)\n",
    "\n",
    "if result.is_success():\n",
    "    P = result.GetSolution(P)\n",
    "    print(\"eig(P) =\" + str(np.linalg.eig(P)[0]))\n",
    "    for i in range(len(A)):\n",
    "        print(\"eig(Pdot\" + str(i) + \") = \" +\n",
    "              str(np.linalg.eig(A[i].transpose().dot(P) + P.dot(A[i]))[0]))\n",
    "else:\n",
    "    print(\"Could not find a common Lyapunov function.\")\n",
    "    print(\"This is expected to occur with some probability:  not all\")\n",
    "    print(\"random sets of stable matrices will have a common Lyapunov\")\n",
    "    print(\"function.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verifying a Lyapunov candidate via SOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import MathematicalProgram, Solve\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewIndeterminates(2, \"x\")\n",
    "f = [-x[0] - 2 * x[1]**2, -x[1] - x[0] * x[1] - 2 * x[1]**3]\n",
    "\n",
    "V = x[0]**2 + 2 * x[1]**2\n",
    "Vdot = V.Jacobian(x).dot(f)\n",
    "\n",
    "prog.AddSosConstraint(-Vdot)\n",
    "\n",
    "result = Solve(prog)\n",
    "assert result.is_success()\n",
    "\n",
    "print(\"Successfully verified Lyapunov candidate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Searching for a Lyapunov function via SOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import MathematicalProgram, Solve, Polynomial, Variables\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewIndeterminates(2, \"x\")\n",
    "f = [-x[0] - 2 * x[1]**2, -x[1] - x[0] * x[1] - 2 * x[1]**3]\n",
    "\n",
    "V = prog.NewSosPolynomial(Variables(x), 2)[0].ToExpression()\n",
    "prog.AddLinearConstraint(V.Substitute({x[0]: 0, x[1]: 0}) == 0)\n",
    "prog.AddLinearConstraint(V.Substitute({x[0]: 1, x[1]: 0}) == 1)\n",
    "Vdot = V.Jacobian(x).dot(f)\n",
    "\n",
    "prog.AddSosConstraint(-Vdot)\n",
    "\n",
    "result = Solve(prog)\n",
    "assert result.is_success()\n",
    "\n",
    "print(\"V = \" + str(\n",
    "    Polynomial(result.GetSolution(V)).RemoveTermsWithSmallCoefficients(1e-5)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Region of attraction for the one-dimensional cubic system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from pydrake.all import Jacobian, MathematicalProgram, Solve\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewIndeterminates(1, \"x\")\n",
    "\n",
    "# Define the dynamics and Lyapunov function.\n",
    "f = -x + x**3\n",
    "V = x.dot(x)\n",
    "Vdot = Jacobian([V], x).dot(f)[0]\n",
    "\n",
    "# Define the Lagrange multiplier.\n",
    "lambda_ = prog.NewSosPolynomial(Variables(x), 2)[0].ToExpression()\n",
    "\n",
    "prog.AddSosConstraint(-Vdot - lambda_*(1 - V))\n",
    "\n",
    "result = Solve(prog)\n",
    "\n",
    "assert result.is_success(), \"Optimization failed\"\n",
    "\n",
    "print(\"Verified that \" + str(V) + \" < 1 is in the region of attraction.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from pydrake.all import Jacobian, MathematicalProgram, Solve\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewIndeterminates(1, \"x\")\n",
    "rho = prog.NewContinuousVariables(1, \"rho\")[0]\n",
    "\n",
    "# Define the dynamics and Lyapunov function.\n",
    "f = -x + x**3\n",
    "V = x.dot(x)\n",
    "Vdot = Jacobian([V], x).dot(f)[0]\n",
    "\n",
    "# Define the Lagrange multiplier.\n",
    "lambda_ = prog.NewContinuousVariables(1, \"lambda\")[0]\n",
    "\n",
    "prog.AddSosConstraint((V - rho) * x.dot(x) - lambda_ * Vdot)\n",
    "prog.AddLinearCost(-rho)\n",
    "\n",
    "result = Solve(prog)\n",
    "\n",
    "assert result.is_success()\n",
    "\n",
    "print(\"Verified that \" + str(V) + \" < \" + str(result.GetSolution(rho)) +\n",
    "      \" is in the region of attraction.\")\n",
    "\n",
    "assert math.fabs(result.GetSolution(rho) - 1) < 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Region of Attraction codes in Drake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import Variable, SymbolicVectorSystem, RegionOfAttraction\n",
    "\n",
    "x = Variable(\"x\")\n",
    "sys = SymbolicVectorSystem(state=[x], dynamics=[-x+x**3])\n",
    "context = sys.CreateDefaultContext()\n",
    "V = RegionOfAttraction(system=sys, context=context)\n",
    "\n",
    "print(\"Verified that \" + str(V) + \" < \" + str(result.GetSolution(rho)) +\n",
    "      \" is in the region of attraction.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Time-reversed van der Pol Oscillator\n",
    "\n",
    "We also use this example in the exercises at the end of the Lyapunov chapter, and work through all of the details of the formulation.  I highly recommend you try it out!\n",
    "\n",
    "N.B. -- we know how to get much larger/tighter (inner) approximations of this RoA.  I will be implementing those ideas in drake's RegionOfAttraction method very soon.  Hopefully the region will be even bigger next time you try it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.all import Variable, SymbolicVectorSystem, RegionOfAttraction, plot_sublevelset_expression\n",
    "from pydrake.examples.van_der_pol import VanDerPolOscillator\n",
    "from underactuated import plot_2d_phase_portrait\n",
    "\n",
    "x1 = Variable(\"x1\")\n",
    "x2 = Variable(\"x2\")\n",
    "sys = SymbolicVectorSystem(state=[x1, x2], dynamics=[-x2, x1 + (x1*x1-1)*x2])\n",
    "context = sys.CreateDefaultContext()\n",
    "V = RegionOfAttraction(system=sys, context=context)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "plot_2d_phase_portrait(sys, (-3, 3), (-3, 3))\n",
    "limit_cycle = VanDerPolOscillator.CalcLimitCycle()\n",
    "plt.plot(limit_cycle[0], limit_cycle[1], color='k', linewidth=3, label='Known ROA boundary')\n",
    "plt.legend(loc=1)\n",
    "plot_sublevelset_expression(ax, V);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimated regions of attraction need not be convex regions (in state space)\n",
    "\n",
    "To demonstrate that, let's make a system with a known, non-convex region of attraction.  We'll do this by taking some interesting potential function $U(x) \\in SOS$ and setting the dynamics to be $\\dot{x} = (U(x)-1) \\frac{\\partial U}{\\partial x}^T$, which has $U(x) <= 1$ as the region of attraction.  \n",
    "\n",
    "Slightly more general is to write $\\dot{x} = (U(x)-1) {\\bf R}(\\theta) \\frac{\\partial U}{\\partial x}^T$, where ${\\bf R}(\\theta) = \\begin{bmatrix} \\cos\\theta & \\sin\\theta \\\\ -\\sin\\theta & \\cos\\theta\\end{bmatrix}$ is the 2D rotation matrix, and $\\theta<\\pi$ is a constant parameter (not a decision variable nor indeterminate), which still has the same region of attraction.\n",
    "\n",
    "**Note: This example is (currently) broken when using Mosek as the SDP solver.**  The estimated region of attraction is actually outside the known region of attraction!  This is due to some subtle numerical issues in the solver.  I've left it in as honest and educational, but we are working to make the solver pre-processing more robust to examples like this.  See [Drake Issue #12876](https://github.com/RobotLocomotion/drake/issues/12876)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pydrake.all import (Variable, Jacobian, SymbolicVectorSystem, \n",
    "                         RegionOfAttraction, RegionOfAttractionOptions, \n",
    "                         plot_sublevelset_expression)\n",
    "\n",
    "# Construct a non-convex 2D level set.\n",
    "x = np.array([Variable(\"x\"), Variable(\"y\")]).reshape((2,))\n",
    "A1 = np.array([[1, 2], [3, 4]])\n",
    "A2 = A1 @ np.array([[-1, 0], [0, 1]])  # mirror about y-axis\n",
    "U = (x.T.dot(A1.T.dot(A1.dot(x)))) * (x.T.dot(A2.T.dot(A2.dot(x))))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "theta = 0.5;\n",
    "R = np.array([[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]])\n",
    "dUdx = U.Jacobian(x)\n",
    "sys = SymbolicVectorSystem(state=x, dynamics=(U-1)* dUdx.T)\n",
    "context = sys.CreateDefaultContext()\n",
    "\n",
    "options = RegionOfAttractionOptions()\n",
    "options.lyapunov_candidate = x.dot(x)\n",
    "options.state_variables = x\n",
    "V = RegionOfAttraction(sys, context, options)\n",
    "plot_sublevelset_expression(ax, V)\n",
    "plot_sublevelset_expression(ax, U, 101, linewidth=3, fill=False);\n",
    "\n",
    "from underactuated import plot_2d_phase_portrait\n",
    "plot_2d_phase_portrait(sys, (-.8, .8), (-.6, .6))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Global stability of the simple pendulum via SOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import MathematicalProgram, Solve, Variables\n",
    "from pydrake.symbolic import Polynomial\n",
    "from pydrake.examples.pendulum import PendulumParams\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "# Declare the \"indeterminates\", x.  These are the variables which define the\n",
    "# polynomials, but are NOT decision variables in the optimization.  We will\n",
    "# add constraints below that must hold FOR ALL x.\n",
    "s = prog.NewIndeterminates(1, \"s\")[0]\n",
    "c = prog.NewIndeterminates(1, \"c\")[0]\n",
    "thetadot = prog.NewIndeterminates(1, \"thetadot\")[0]\n",
    "# TODO(russt): bind the sugar methods so I can write\n",
    "#  x = prog.NewIndeterminates([\"s\", \"c\", \"thetadot\"])\n",
    "x = np.array([s, c, thetadot])\n",
    "\n",
    "# Write out the dynamics in terms of sin(theta), cos(theta), and thetadot\n",
    "p = PendulumParams()\n",
    "f = [\n",
    "    c * thetadot, -s * thetadot,\n",
    "    (-p.damping() * thetadot - p.mass() * p.gravity() * p.length() * s) /\n",
    "    (p.mass() * p.length() * p.length())\n",
    "]\n",
    "\n",
    "# The fixed-point in this coordinate (because cos(0)=1).\n",
    "x0 = np.array([0, 1, 0])\n",
    "\n",
    "# Construct a polynomial V that contains all monomials with s,c,thetadot up\n",
    "# to degree 2.\n",
    "deg_V = 2\n",
    "V = prog.NewFreePolynomial(Variables(x), deg_V).ToExpression()\n",
    "\n",
    "# Add a constraint to enforce that V is strictly positive away from x0.\n",
    "# (Note that because our coordinate system is sine and cosine, V is also zero\n",
    "# at theta=2pi, etc).\n",
    "eps = 1e-4\n",
    "constraint1 = prog.AddSosConstraint(V - eps * (x - x0).dot(x - x0))\n",
    "\n",
    "# Construct the polynomial which is the time derivative of V.\n",
    "Vdot = V.Jacobian(x).dot(f)\n",
    "\n",
    "# Construct a polynomial L representing the \"Lagrange multiplier\".\n",
    "deg_L = 2\n",
    "L = prog.NewFreePolynomial(Variables(x), deg_L).ToExpression()\n",
    "\n",
    "# Add a constraint that Vdot is strictly negative away from x0 (but make an\n",
    "# exception for the upright fixed point by multipling by s^2).\n",
    "constraint2 = prog.AddSosConstraint(-Vdot - L * (s**2 + c**2 - 1) - eps *\n",
    "                                    (x - x0).dot(x - x0) * s**2)\n",
    "\n",
    "# Add V(0) = 0 constraint\n",
    "constraint3 = prog.AddLinearConstraint(\n",
    "    V.Substitute({\n",
    "        s: 0,\n",
    "        c: 1,\n",
    "        thetadot: 0\n",
    "    }) == 0)\n",
    "\n",
    "# Add V(theta=pi) = mgl, just to set the scale.\n",
    "constraint4 = prog.AddLinearConstraint(\n",
    "    V.Substitute({\n",
    "        s: 1,\n",
    "        c: 0,\n",
    "        thetadot: 0\n",
    "    }) == p.mass() * p.gravity() * p.length())\n",
    "\n",
    "# Call the solver.\n",
    "result = Solve(prog)\n",
    "assert result.is_success()\n",
    "\n",
    "# Note that I've added mgl to the potential energy (relative to the textbook),\n",
    "# so that it would be non-negative... like the Lyapunov function.\n",
    "mgl = p.mass() * p.gravity() * p.length()\n",
    "print(\"Mechanical Energy = \")\n",
    "print(.5 * p.mass() * p.length()**2 * thetadot**2 + mgl * (1 - c))\n",
    "\n",
    "print(\"V =\")\n",
    "Vsol = Polynomial(result.GetSolution(V))\n",
    "print(Vsol.RemoveTermsWithSmallCoefficients(1e-6))\n",
    "\n",
    "# Plot the results as contour plots.\n",
    "nq = 151\n",
    "nqd = 151\n",
    "q = np.linspace(-2 * np.pi, 2 * np.pi, nq)\n",
    "qd = np.linspace(-2 * mgl, 2 * mgl, nqd)\n",
    "Q, QD = np.meshgrid(q, qd)\n",
    "Energy = .5 * p.mass() * p.length()**2 * QD**2 + mgl * (1 - np.cos(Q))\n",
    "Vplot = Q.copy()\n",
    "env = {s: 0., c: 1., thetadot: 0}\n",
    "for i in range(nq):\n",
    "    for j in range(nqd):\n",
    "        env[s] = np.sin(Q[i, j])\n",
    "        env[c] = np.cos(Q[i, j])\n",
    "        env[thetadot] = QD[i, j]\n",
    "        Vplot[i, j] = Vsol.Evaluate(env)\n",
    "\n",
    "# plt.rc(\"text\", usetex=True)\n",
    "fig, ax = plt.subplots()\n",
    "ax.contour(Q, QD, Vplot)\n",
    "ax.contour(Q, QD, Energy, alpha=0.5, linestyles=\"dashed\")\n",
    "ax.set_xlabel(\"theta\")\n",
    "ax.set_ylabel(\"thetadot\")\n",
    "ax.set_title(\"V (solid) and Mechanical Energy (dashed)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
